package org.spring.aicloud.controller;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.volcengine.ark.runtime.model.images.generation.ImagesResponse;
import io.minio.errors.*;
import jakarta.annotation.Resource;
import org.spring.aicloud.entity.Answer;
import okhttp3.Dispatcher;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import org.spring.aicloud.entity.enums.AIModelEnum;
import org.spring.aicloud.entity.enums.AITypeEnum;
import org.spring.aicloud.service.IAnswerService;
import org.spring.aicloud.utils.MinIoUtil;
import org.spring.aicloud.utils.ResponseEntity;
import org.spring.aicloud.utils.SecurityUtil;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * @author yexiebao
 * @date 2025/6/1
 * 字节大模型
 */
@RestController
@RequestMapping("/doubao")
public class DoubaoController {
    @Value("${doubao.api-key}")
    private String apiKey;
    @Value("${doubao.url}")
    private String url;
    @Resource
    private IAnswerService  answerService;
    @Resource
    private MinIoUtil minIoUtil;

    /**
     * 对话
     * @param question
     * @return
     */
    @RequestMapping("/chat")
    public ResponseEntity chat(String question) {
        if(!StringUtils.hasLength(question)){
            return ResponseEntity.error("请先输入内容");
        }
        //1.传入api-key,调用url地址,创建服务
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(url).apiKey(apiKey).build();
        //2.构建问题对象
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(question).build();
        messages.add(userMessage);
        //3.构建请求体
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model("ep-20250601215039-zzvq9")
                .messages(messages)
                .build();
        //4.调用服务,拿到结果
        StringBuilder output = new StringBuilder();
        service.createChatCompletion(chatCompletionRequest)
                .getChoices()
                .forEach(choice -> output.append(choice.getMessage().getContent()).append("\n"));
        String result = output.toString();
        //保存答案到数据库
        Answer answer = new Answer();
        answer.setTitle(question);
        answer.setContent(result);
        answer.setModel(AIModelEnum.DOUBAO.getCode());
        answer.setType(AITypeEnum.CHAT.getCode());
        answer.setUid(SecurityUtil.getCurrentUser().getUid());
        boolean saveResult = answerService.save(answer);
        if(!saveResult){
            return ResponseEntity.error("结果保存失败");
        }
        return ResponseEntity.ok(result);
    }
    /**
     * 图片生成
     * @param question
     * @return
     */
    @RequestMapping("/draw")
    public ResponseEntity draw(String question) throws IOException {
        if(!StringUtils.hasLength(question)){
            return ResponseEntity.error("请先输入内容");
        }
        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).apiKey(apiKey).build();

        GenerateImagesRequest generateRequest = GenerateImagesRequest.builder()
                .model("doubao-seedream-3-0-t2i-250415")
                .prompt(question)
                .build();

        ImagesResponse imagesResponse = service.generateImages(generateRequest);
        String resultUrl = imagesResponse.getData().get(0).getUrl();
        service.shutdownExecutor();
        URL downloadUrl = new URL(resultUrl);
        //上传到minIO
        String imgUrl;
        try(InputStream inputStream = downloadUrl.openStream()) {
            imgUrl = minIoUtil.upload("doubao-draw-" + UUID.randomUUID().toString().replace("-", ""), inputStream, "image/png");
        }catch (ServerException | InsufficientDataException | ErrorResponseException |
                NoSuchAlgorithmException | InvalidKeyException | InvalidResponseException |
                XmlParserException | InternalException | IOException e) {
            // 抛出运行时异常
            throw new RuntimeException("上传图片到 MinIO 失败", e);
        }
        //保存答案到数据库
        Answer answer = new Answer();
        answer.setTitle(question);
        answer.setContent(imgUrl);
        answer.setModel(AIModelEnum.DOUBAO.getCode());
        answer.setType(AITypeEnum.DRAW.getCode());
        answer.setUid(SecurityUtil.getCurrentUser().getUid());
        boolean saveResult = answerService.save(answer);
        if(!saveResult){
            return ResponseEntity.error("结果保存失败");
        }
        return ResponseEntity.ok(imgUrl);
    }
    /**
     * 获取对话历史
     */
    @RequestMapping("/getchatlist")
    public ResponseEntity getChatList() {
        Long uid = SecurityUtil.getCurrentUser().getUid();
        QueryWrapper<Answer> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("uid",uid);
        queryWrapper.eq("type", AITypeEnum.CHAT.getCode());
        queryWrapper.eq("model", AIModelEnum.DOUBAO.getCode());
        queryWrapper.orderByDesc("aid");
        return ResponseEntity.ok(answerService.list(queryWrapper));
    }
    /**
     * 获取绘画历史
     */
    @RequestMapping("/getdrawlist")
    public ResponseEntity getDrawList() {
        Long uid =SecurityUtil.getCurrentUser().getUid();
        QueryWrapper<Answer> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("uid",uid);
        queryWrapper.eq("type", AITypeEnum.DRAW.getCode());
        queryWrapper.eq("model",AIModelEnum.DOUBAO.getCode());
        queryWrapper.orderByDesc("aid");
        return ResponseEntity.ok(answerService.list(queryWrapper));
    }
}
